package cn.com.guage.transtation;

import com.mysql.cj.jdbc.MysqlXADataSource;

import javax.sql.XAConnection;
import javax.transaction.xa.XAResource;
import javax.transaction.xa.Xid;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.util.UUID;

public class DistributedTransaction2PC {

    public static void main(String[] args) throws Exception {
        // 配置XA数据源
        MysqlXADataSource dsBank1 = createXADataSource("jdbc:mysql://localhost:3306/2pc-db1", "root", "123456");
        MysqlXADataSource dsBank2 = createXADataSource("jdbc:mysql://localhost:3306/2pc-db2", "root", "123456");

        XAConnection xaConnBank1 = null;
        XAConnection xaConnBank2 = null;

        try {
            // 获取XA连接和资源
            xaConnBank1 = dsBank1.getXAConnection();
            XAResource xaResBank1 = xaConnBank1.getXAResource();

            xaConnBank2 = dsBank2.getXAConnection();
            XAResource xaResBank2 = xaConnBank2.getXAResource();

            // 生成全局事务ID
            byte[] gtrid = UUID.randomUUID().toString().getBytes();
            int formatId = 1;

            // 创建事务分支ID
            Xid xidBank1 = new MyXid(gtrid, new byte[]{0x01}, formatId);
            Xid xidBank2 = new MyXid(gtrid, new byte[]{0x02}, formatId);

            // 分支事务1：扣除张三的金额
            xaResBank1.start(xidBank1, XAResource.TMNOFLAGS);
            Connection connBank1 = xaConnBank1.getConnection();
            PreparedStatement psBank1 = connBank1.prepareStatement(
                    "UPDATE account SET balance = balance - 1000 WHERE name = '张三'");
            psBank1.executeUpdate();
            xaResBank1.end(xidBank1, XAResource.TMSUCCESS);

            // 分支事务2：增加李四的金额
            xaResBank2.start(xidBank2, XAResource.TMNOFLAGS);
            Connection connBank2 = xaConnBank2.getConnection();
            PreparedStatement psBank2 = connBank2.prepareStatement(
                    "UPDATE account SET balance = balance + 1000 WHERE name = '李四'");
            psBank2.executeUpdate();
            xaResBank2.end(xidBank2, XAResource.TMSUCCESS);

            // 两阶段提交协议
            int prepareBank1 = xaResBank1.prepare(xidBank1);
            int prepareBank2 = xaResBank2.prepare(xidBank2);

            if (prepareBank1 == XAResource.XA_OK && prepareBank2 == XAResource.XA_OK) {
                xaResBank1.commit(xidBank1, false);
                xaResBank2.commit(xidBank2, false);
                System.out.println("事务提交成功！");
            } else {
                xaResBank1.rollback(xidBank1);
                xaResBank2.rollback(xidBank2);
                System.out.println("事务回滚！");
            }
        } catch (Exception e) {
            // 异常处理，回滚事务
            if (xaConnBank1 != null) xaConnBank1.getXAResource().rollback(new MyXid(new byte[0], new byte[0], 0));
            if (xaConnBank2 != null) xaConnBank2.getXAResource().rollback(new MyXid(new byte[0], new byte[0], 0));
            e.printStackTrace();
            System.err.println("事务执行失败，已回滚！");
        } finally {
            if (xaConnBank1 != null) xaConnBank1.close();
            if (xaConnBank2 != null) xaConnBank2.close();
        }
    }

    private static MysqlXADataSource createXADataSource(String url, String user, String password) {
        MysqlXADataSource ds = new MysqlXADataSource();
        ds.setUrl(url);
        ds.setUser(user);
        ds.setPassword(password);
        return ds;
    }

    static class MyXid implements Xid {
        private final byte[] gtrid;
        private final byte[] bqual;
        private final int formatId;

        public MyXid(byte[] gtrid, byte[] bqual, int formatId) {
            this.gtrid = gtrid;
            this.bqual = bqual;
            this.formatId = formatId;
        }

        @Override
        public byte[] getGlobalTransactionId() {
            return gtrid;
        }

        @Override
        public byte[] getBranchQualifier() {
            return bqual;
        }

        @Override
        public int getFormatId() {
            return formatId;
        }
    }
}